import os
import sys
import shutil
import re


# 根据文件名中的关键字，自动重命名文件


SRC_DIR = sys.argv[1]

print("SRC_DIR", SRC_DIR)

KEYWORDS = ["工业投资运行预计", "5亿元以上项目", "能源外工业投资预计"]
TARGET_NAMES = ["gy.xlsx", "51.xlsx", "ne.xlsx"]


def get_files(dir, ext, recursive=False):
    """
    获取指定目录下指定扩展名的文件列表，默认不递归
    :param dir: 目录路径
    :param ext: 扩展名，带点 `.xlsx`
    :param recursive: 是否递归查找
    :return: 文件列表
    """
    all_files = []

    if recursive:
        # 遍历指定目录及其子目录下的所有文件
        for root, _, files in os.walk(dir):
            for file in files:
                # 使用os.path.join拼接完整的文件路径
                all_files.append(os.path.join(root, file))
    else:
        # 只取一层文件夹的文件
        all_files = map(lambda x: os.path.join(dir, x), os.listdir(dir))

    # 排除扩展名不匹配的文件，还有 `~$` 开头的临时文件
    valid_ext_files = filter(
        lambda x: x.endswith(ext) and not os.path.basename(x).startswith(".~"),
        all_files,
    )

    return list(valid_ext_files)


def try_file_path(file_path):
    """
    测试是否存在文件，如果存在，返回带括号的文件名
    :param file_path: 文件路径
    :return: 新文件名
    """
    tmp_path = file_path
    dir_and_name, ext = os.path.splitext(tmp_path)
    i = 1
    while os.path.exists(tmp_path):
        tmp_path = f"{dir_and_name}({i}){ext}"
        i += 1
    return tmp_path


def check_rename_files(files):
    for f in files:
        for k in KEYWORDS:
            filename = os.path.basename(f)
            mat = re.search(k, filename)
            if mat is not None:
                print("Move", f, k)
                ind = KEYWORDS.index(k)
                target = try_file_path(f"{TARGET_NAMES[ind]}")
                shutil.copy(f, target)


def main():
    files = get_files(SRC_DIR, ".xlsx", True)
    print(files)
    check_rename_files(files)

    print("Done")


if __name__ == "__main__":
    main()
